/*++

Copyright (c) 2023 Shorthorn Project

Module Name:

    main.c

Abstract:

    This module implements COM Apartment functions APIs

Author:

    Skulltrail 12-October-2023

Revision History:

--*/

#include "main.h"

WINE_DEFAULT_DEBUG_CHANNEL(ole32);

BOOLEAN g_cMTAInits;  // Global variables within ole32 itself
DWORD gdwMainThreadId;
static struct apartment *mta;
static struct list apts = LIST_INIT(apts);

struct mta_cookie
{
    struct list entry;
};

static CRITICAL_SECTION apt_cs;
static CRITICAL_SECTION_DEBUG apt_cs_debug =
{
    0, 0, &apt_cs,
    { &apt_cs_debug.ProcessLocksList, &apt_cs_debug.ProcessLocksList },
      0, 0, { (DWORD_PTR)(__FILE__ ": apt_cs") }
};
static CRITICAL_SECTION apt_cs = { &apt_cs_debug, -1, 0, 0, 0, 0 };

static DWORD apartment_addref(struct apartment *apt)
{
    DWORD refs = InterlockedIncrement(&apt->refs);
    TRACE("%s: before = %ld\n", wine_dbgstr_longlong(apt->oxid), refs - 1);
    return refs;
}

HRESULT WINAPI CoGetApartmentType(APTTYPE *pAptType, APTTYPEQUALIFIER *pAptQualifier)
/* 
    win32 - March 7 2021
	An important OLE function introduced in Windows 7, but should be usable in NT 4 and above.
	Modified to better support wrappers
*/
{
  IUnknown ContextToken;
  HRESULT Status; 
  SOleTlsData* ReservedForOle; 
  DWORD CurrentThreadId; 
  DWORD OleFlags; 
  APTTYPE AptTypeTemp; 
  APTTYPEQUALIFIER AptQualifierTemp; 
  TEB *CurrentThreadInfo;
  
  Status = S_OK;
  if ( !pAptType || !pAptQualifier )
    return E_INVALIDARG;
  *pAptType = APTTYPE_CURRENT;
  *pAptQualifier = APTTYPEQUALIFIER_NONE;
  CurrentThreadInfo = NtCurrentTeb();
  ReservedForOle = (SOleTlsData*) CurrentThreadInfo->ReservedForOle;
  CurrentThreadId = GetCurrentThreadId();

  if ( !ReservedForOle )
  {
	// Try to see if COM is initialized
    if ( CoGetContextToken((ULONG_PTR*)&ContextToken) == 0)
      goto ImplicitMTA;
    return CO_E_NOTINITIALIZED;
  }
  OleFlags = ReservedForOle->dwFlags;
  if ( (OleFlags & 0x800) == 0 )
  {
    if ( (OleFlags & 0x80) != 0 )
    {
      AptTypeTemp = APTTYPE_STA;
      if (FALSE) // not supported yet if ( CurrentThreadId == gdwMainThreadId )
        AptTypeTemp = APTTYPE_MAINSTA;
      *pAptType = AptTypeTemp;
      if ( (OleFlags & 0x400000) != 0 )
        *pAptQualifier = APTTYPEQUALIFIER_APPLICATION_STA;
   
      if ( (OleFlags & 0x40000000) != 0 )
        *pAptQualifier = APTTYPEQUALIFIER_RESERVED_1;
      return Status;
    }
    if ( (OleFlags & 0x1100) != 0 )
    {
      *pAptType = APTTYPE_MTA;
      return Status;
    }
    if ( CoGetContextToken((ULONG_PTR*)&ContextToken) == 0 )
    {
ImplicitMTA:
      *pAptType = APTTYPE_MTA;
      *pAptQualifier = APTTYPEQUALIFIER_IMPLICIT_MTA;
      return Status;
    }
    return CO_E_NOTINITIALIZED;
  }
  AptQualifierTemp = APTTYPEQUALIFIER_NA_ON_MTA;
  *pAptType = APTTYPE_NA;
  if ( (OleFlags & 0x80) != 0 )
  {
    AptQualifierTemp = APTTYPEQUALIFIER_NA_ON_STA;
    if (FALSE) // not supported yet if ( CurrentThreadId == gdwMainThreadId )
      AptQualifierTemp = APTTYPEQUALIFIER_NA_ON_MAINSTA;
    *pAptQualifier = AptQualifierTemp;
	return Status;
  }
  if ( (OleFlags & 0x1100) != 0 )
  {
    *pAptQualifier = AptQualifierTemp;
    return Status;
  }
  if ( TRUE ) // TODO: make replacement for this.
  {
    *pAptQualifier = APTTYPEQUALIFIER_NA_ON_IMPLICIT_MTA;
    return Status;
  }
  return E_FAIL;
}

HRESULT ensure_mta(void)
{
    struct apartment *apt;
    struct tlsdata *data;
    HRESULT hr;

    if (FAILED(hr = com_get_tlsdata(&data)))
        return hr;
    if ((apt = data->apt) && (data->implicit_mta_cookie || apt->multi_threaded))
        return S_OK;

    EnterCriticalSection(&apt_cs);
    if (apt || mta)
        hr = apartment_increment_mta_usage(&data->implicit_mta_cookie);
    else
        hr = CO_E_NOTINITIALIZED;
    LeaveCriticalSection(&apt_cs);

    if (FAILED(hr))
    {
        ERR("Failed, hr %#lx.\n", hr);
        return hr;
    }
    return S_OK;
}

/* Creates new apartment for given model */
static struct apartment *apartment_construct(DWORD model)
{
    struct apartment *apt;

    TRACE("creating new apartment, model %ld\n", model);

    apt = calloc(1, sizeof(*apt));
    apt->tid = GetCurrentThreadId();

    list_init(&apt->proxies);
    list_init(&apt->stubmgrs);
    list_init(&apt->loaded_dlls);
    list_init(&apt->usage_cookies);
    apt->ipidc = 0;
    apt->refs = 1;
    apt->remunk_exported = FALSE;
    apt->oidc = 1;
    InitializeCriticalSectionEx(&apt->cs, 0, RTL_CRITICAL_SECTION_FLAG_FORCE_DEBUG_INFO);
    apt->cs.DebugInfo->Spare[0] = (DWORD_PTR)(__FILE__ ": apartment");

    apt->multi_threaded = !(model & COINIT_APARTMENTTHREADED);

    if (apt->multi_threaded)
    {
        /* FIXME: should be randomly generated by in an RPC call to rpcss */
        apt->oxid = ((OXID)GetCurrentProcessId() << 32) | 0xcafe;
    }
    else
    {
        /* FIXME: should be randomly generated by in an RPC call to rpcss */
        apt->oxid = ((OXID)GetCurrentProcessId() << 32) | GetCurrentThreadId();
    }

    TRACE("Created apartment on OXID %s\n", wine_dbgstr_longlong(apt->oxid));

    list_add_head(&apts, &apt->entry);

    return apt;
}

HRESULT apartment_increment_mta_usage(CO_MTA_USAGE_COOKIE *cookie)
{
    struct mta_cookie *mta_cookie;

    *cookie = NULL;

    if (!(mta_cookie = malloc(sizeof(*mta_cookie))))
        return E_OUTOFMEMORY;

    EnterCriticalSection(&apt_cs);

    if (mta)
        apartment_addref(mta);
    else
        mta = apartment_construct(COINIT_MULTITHREADED);
    list_add_head(&mta->usage_cookies, &mta_cookie->entry);

    LeaveCriticalSection(&apt_cs);

    *cookie = (CO_MTA_USAGE_COOKIE)mta_cookie;

    return S_OK;
}